import torch
from torch import nn
from ops.sigmoid_module import Sigmoid

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

    def forward(self, x):
        return Sigmoid.apply(x)

device = 'cuda:0'
input = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.float32,requires_grad=True).to(device)

net = Net().to(device)
output = net(input)
loss = output.sum()
loss.backward()
pass
